import math

import tensorflow as tf


class MLP(tf.keras.Model):
    def __init__(self):
        super().__init__()
        feature_num = 30  # 输入数据特征数
        start_node = 2 ** (math.ceil(math.log(feature_num / 2,
                                              2)))  # 第一层节点数量

        self.dense0 = tf.keras.layers.Dense(units=256, activation=tf.nn.relu)
        self.dense1 = tf.keras.layers.Dense(units=128, activation=tf.nn.relu)
        self.dense2 = tf.keras.layers.Dense(units=64, activation=tf.nn.relu)
        self.dense3 = tf.keras.layers.Dense(units=32, activation=tf.nn.relu)
        self.dense4 = tf.keras.layers.Dense(units=16, activation=tf.nn.relu)
        self.dense5 = tf.keras.layers.Dense(units=8, activation=tf.nn.relu)
        # self.dense6 = tf.keras.layers.Dense(units=4, activation=tf.nn.relu)
        self.dense7 = tf.keras.layers.Dense(units=2)
        self.dropout1 = tf.keras.layers.Dropout(0.5)
        self.dropout2 = tf.keras.layers.Dropout(0.3)

    def call(self, inputs):
        x = self.dense0(inputs)
        x = self.dense1(x)
        x = self.dropout1(x)
        x = self.dense2(x)
        x = self.dropout1(x)
        x = self.dense3(x)
        x = self.dropout1(x)
        x = self.dense4(x)
        # x = self.dropout2(x)
        x = self.dense5(x)
        # x = self.dense6(x)
        x = self.dense7(x)  # [batch,2]是两个类别的概率
        output = tf.nn.softmax(x)
        return output
